package gti310.tp4;

import java.util.ArrayList;

/**
 * The Main class is where the different functions are called to either encode a
 * PPM file to the Squeeze-Light format or to decode a Squeeze-Ligth image into
 * PPM format. It is the implementation of the simplified JPEG block diagrams.
 * 
 * @author Fran�ois Caron
 */
public class Main {

	/*
	 * The entire application assumes that the blocks are 8x8 squares.
	 */
	public static final int BLOCK_SIZE = 8;

	/*
	 * The number of dimensions in the color spaces.
	 */
	public static final int COLOR_SPACE_SIZE = 3;

	/*
	 * The RGB color space.
	 */
	public static final int R = 0;
	public static final int G = 1;
	public static final int B = 2;

	/*
	 * The YUV color space.
	 */
	public static final int Y = 0;
	public static final int U = 1;
	public static final int V = 2;

	/**
	 * The application's entry point.
	 * 
	 * @param args
	 */
	public static void main(String[] args) {
		System.out.println("Squeeze Light Media Codec !");
		int qualityFactor;

		try {
			qualityFactor = Integer.parseInt(args[2]);
			code(PPMReaderWriter.readPPMFile(args[0]), args[1], qualityFactor);
		} catch (ArrayIndexOutOfBoundsException e) {
			// No Quality factor so switching to decode mode
			decode(SZLReaderWriter.readSZLFile(args[0]), args[1]);
		}
		System.out.println("Done !");

	}

	private static void decode(int[] imageHeader, String destination) {
		int[][][] imageRBG = null;
		int numberOfBlock = imageHeader[0] * imageHeader[1]
				/ (Main.BLOCK_SIZE * Main.BLOCK_SIZE);
		int vectorLength = Main.BLOCK_SIZE * Main.BLOCK_SIZE;
		int[] YDCvector = new int[numberOfBlock];
		int[] UDCvector = new int[numberOfBlock];
		int[] VDCvector = new int[numberOfBlock];
		int[][][] imageYUV = new int[imageHeader[2]][imageHeader[0]][imageHeader[1]];
		double[][] YColor = new double[imageHeader[0]][imageHeader[1]];
		double[][] UColor = new double[imageHeader[0]][imageHeader[1]];
		double[][] VColor = new double[imageHeader[0]][imageHeader[1]];
		ArrayList<double[][]> YiDCTList = new ArrayList<double[][]>();
		ArrayList<double[][]> UiDCTList = new ArrayList<double[][]>();
		ArrayList<double[][]> ViDCTList = new ArrayList<double[][]>();
		ArrayList<int[]> YACDCList = new ArrayList<int[]>();
		ArrayList<int[]> UACDCList = new ArrayList<int[]>();
		ArrayList<int[]> VACDCList = new ArrayList<int[]>();

		// Get AC-DC newest album
		// Y (DC)
		// O(n)
		int lastDC = 0;
		for (int i = 0; i < numberOfBlock; i++) {
			YDCvector[i] = Entropy.readDC() + lastDC;
			lastDC = YDCvector[i];
		}
		// U (DC)
		// O(n)
		lastDC = 0;
		for (int i = 0; i < numberOfBlock; i++) {
			UDCvector[i] = Entropy.readDC() + lastDC;
			lastDC = UDCvector[i];
		}
		// V (DC)
		// O(n)
		lastDC = 0;
		for (int i = 0; i < numberOfBlock; i++) {
			VDCvector[i] = Entropy.readDC() + lastDC;
			lastDC = VDCvector[i];
		}

		// Y (AC)
		// O(n^2)
		int[] AC = null;
		int[] acVector;
		int padding;
		for (int i = 0; i < numberOfBlock; i++) {

			AC = null;
			acVector = new int[vectorLength];
			padding = 1; // The +1 is to make sure everything is shifted
							// from one place to leave a space for DC
			do {
				AC = Entropy.readAC();
				if (AC[1] == 0) {// Zeros till the end
					for (; padding < acVector.length; padding++) {
						acVector[padding] = 0;
					}
				} else if (AC[0] != 0) {// Some spaced to fill with zeros
					for (int j = 0; j < AC[0]; j++) {
						acVector[padding] = 0;
						padding++;
					}
					acVector[padding] = AC[1];// Add the value
					padding++;
				} else {// No zeros
					acVector[padding] = AC[1];
					padding++;
				}

			} while (padding != vectorLength);
			acVector[0] = YDCvector[i]; // Add DC to the AC vector
			YACDCList.add(acVector);

		}
		// U (AC)
		// O(n^2)
		for (int i = 0; i < numberOfBlock; i++) {

			AC = null;
			acVector = new int[vectorLength];
			padding = 1; // The +1 is to make sure everything is shifted
							// from one place to leave a space for DC
			do {
				AC = Entropy.readAC();
				if (AC[1] == 0) {// Zeros till the end
					for (; padding < acVector.length; padding++) {
						acVector[padding] = 0;
					}
				} else if (AC[0] != 0) {// Some spaced to fill with zeros
					for (int j = 0; j < AC[0]; j++) {
						acVector[padding] = 0;
						padding++;
					}
					acVector[padding] = AC[1];// Add the value
					padding++;
				} else {// No zeros
					acVector[padding] = AC[1];
					padding++;
				}

			} while (padding != vectorLength);
			acVector[0] = UDCvector[i]; // Add DC to the AC vector
			UACDCList.add(acVector);

		}
		// V (AC)
		// O(n^2)
		for (int i = 0; i < numberOfBlock; i++) {

			AC = null;
			acVector = new int[vectorLength];
			padding = 1; // The +1 is to make sure everything is shifted
							// from one place to leave a space for DC
			do {
				AC = Entropy.readAC();
				if (AC[1] == 0) {// Zeros till the end
					for (; padding < acVector.length; padding++) {
						acVector[padding] = 0;
					}
				} else if (AC[0] != 0) {// Some spaced to fill with zeros
					for (int j = 0; j < AC[0]; j++) {
						acVector[padding] = 0;
						padding++;
					}
					acVector[padding] = AC[1];// Add the value
					padding++;
				} else {// No zeros
					acVector[padding] = AC[1];
					padding++;
				}

			} while (padding != vectorLength);
			acVector[0] = VDCvector[i]; // Add DC to the AC vector
			VACDCList.add(acVector);

		}

		// iZigZag, iQuantification, iDCT
		// Y
		// O(n)
		for (int i = 0; i < YACDCList.size(); i++) {
			YiDCTList.add(DCT.deTransform(Quantification.deQuantifierY(
					ZigZag.list2Matrix(YACDCList.get(i), Main.BLOCK_SIZE),
					imageHeader[3])));
		}
		// U
		// O(n)
		for (int i = 0; i < UACDCList.size(); i++) {
			UiDCTList.add(DCT.deTransform(Quantification.deQuantifierUV(
					ZigZag.list2Matrix(UACDCList.get(i), Main.BLOCK_SIZE),
					imageHeader[3])));
		}
		// V
		// O(n)
		for (int i = 0; i < VACDCList.size(); i++) {
			ViDCTList.add(DCT.deTransform(Quantification.deQuantifierUV(
					ZigZag.list2Matrix(VACDCList.get(i), Main.BLOCK_SIZE),
					imageHeader[3])));
		}

		// Un-Chop Chop the 8x8 samples
		// O(n^3)
		int g = 0; // used to select all block in the list
		for (int k = 0; k < imageYUV[0][0].length; k = k + Main.BLOCK_SIZE) {

			for (int i = 0; i < imageYUV[0].length; i++) {
				if (i % Main.BLOCK_SIZE == 0) {
					YColor = YiDCTList.get((i / Main.BLOCK_SIZE) + g);
					UColor = UiDCTList.get((i / Main.BLOCK_SIZE) + g);
					VColor = ViDCTList.get((i / Main.BLOCK_SIZE) + g);
				}
				for (int j = 0; j < Main.BLOCK_SIZE; j++) {
					imageYUV[Main.Y][i][j + k] = (int) Math.round(YColor[i
							% Main.BLOCK_SIZE][j]);
					imageYUV[Main.U][i][j + k] = (int) Math.round(UColor[i
							% Main.BLOCK_SIZE][j]);
					imageYUV[Main.V][i][j + k] = (int) Math.round(VColor[i
							% Main.BLOCK_SIZE][j]);
				}
			}
			g = g + imageYUV[0].length / Main.BLOCK_SIZE;
		}

		imageRBG = ColorConvert.yulrgb(imageYUV);
		PPMReaderWriter.writePPMFile(destination, imageRBG);
	}

	private static void code(int[][][] image, String destination,
			int qualityFactor) {
		int[][][] imageYUV;
		double[][] Yblock = new double[Main.BLOCK_SIZE][Main.BLOCK_SIZE];
		double[][] Ublock = new double[Main.BLOCK_SIZE][Main.BLOCK_SIZE];
		double[][] Vblock = new double[Main.BLOCK_SIZE][Main.BLOCK_SIZE];
		ArrayList<double[][]> YDCTList = new ArrayList<double[][]>();
		ArrayList<double[][]> UDCTList = new ArrayList<double[][]>();
		ArrayList<double[][]> VDCTList = new ArrayList<double[][]>();
		ArrayList<int[]> YACList = new ArrayList<int[]>();
		ArrayList<int[]> UACList = new ArrayList<int[]>();
		ArrayList<int[]> VACList = new ArrayList<int[]>();

		imageYUV = ColorConvert.rgbyul(image);

		int blockHeight = 0;
		int blockWidth = 0;

		int width = imageYUV[0].length;
		int height = imageYUV[0][0].length;

		// Chop Chop to send 8x8 array to DCT
		// O(n^3)
		for (int k = 0; k < height; k = k + Main.BLOCK_SIZE) {
			for (int i = 0; i < width; i++) {
				for (int j = 0; j < height; j++) {

					Yblock[blockHeight][blockWidth] = imageYUV[Main.Y][i][j + k];
					Ublock[blockHeight][blockWidth] = imageYUV[Main.U][i][j + k];
					Vblock[blockHeight][blockWidth] = imageYUV[Main.V][i][j + k];

					if (blockWidth == Main.BLOCK_SIZE - 1
							&& blockHeight == Main.BLOCK_SIZE - 1) { // Block found
						// DCT those blocks and add the result to their list
						YDCTList.add(DCT.transform(Yblock));
						UDCTList.add(DCT.transform(Ublock));
						VDCTList.add(DCT.transform(Vblock));

						blockWidth = -1;// to compensate for the blockWidth++
										// ahead
						blockHeight = 0;
						break;
					}

					if (blockHeight == Main.BLOCK_SIZE - 1) {// Column filled,
																// jump
																// to next
																// column
						blockHeight = 0;
						break;
					} else {
						blockHeight++;
					}

				}

				blockWidth++;
			}
		}
		int[] zigzag;
		int lastDC = 0;

		// Quantification, Zigzag, DC(DPCM) and while we're there, we
		// stock the AC(runLength, value) pair.
		// Y
		// O(n^2)
		for (int i = 0; i < YDCTList.size(); i++) {
			if (i == 0) {
				lastDC = 0;
			}
			zigzag = ZigZag.matrix2List(
					Quantification.quantifierY(YDCTList.get(i), qualityFactor),
					Main.BLOCK_SIZE);
			Entropy.writeDC(zigzag[0] - lastDC); // DPCM
			lastDC = zigzag[0];

			// AC
			int runLength = 0;
			for (int j = 1; j < zigzag.length; j++) {// j=1 to skip the DC
				if (zigzag[j] != 0 || j == zigzag.length - 1) {
					YACList.add(new int[] { runLength, zigzag[j] });
					runLength = 0;
				} else {
					runLength++;
				}
			}

		}

		// U
		// O(n^2)
		for (int i = 0; i < UDCTList.size(); i++) {
			if (i == 0) {
				lastDC = 0;
			}
			zigzag = ZigZag
					.matrix2List(Quantification.quantifierUV(UDCTList.get(i),
							qualityFactor), Main.BLOCK_SIZE);
			Entropy.writeDC(zigzag[0] - lastDC); // DPCM
			lastDC = zigzag[0];

			// AC
			int runLength = 0;
			for (int j = 1; j < zigzag.length; j++) {// j=1 to skip the DC
				if (zigzag[j] != 0 || j == zigzag.length - 1) {
					UACList.add(new int[] { runLength, zigzag[j] });
					runLength = 0;
				} else {
					runLength++;
				}
			}

		}
		
		// V
		// O(n^2)
		for (int i = 0; i < VDCTList.size(); i++) {
			if (i == 0) {
				lastDC = 0;
			}
			zigzag = ZigZag
					.matrix2List(Quantification.quantifierUV(VDCTList.get(i),
							qualityFactor), Main.BLOCK_SIZE);
			Entropy.writeDC(zigzag[0] - lastDC); // DPCM
			lastDC = zigzag[0];

			// AC
			int runLength = 0;
			for (int j = 1; j < zigzag.length; j++) {// j=1 to skip the DC
				if (zigzag[j] != 0 || j == zigzag.length - 1) {
					VACList.add(new int[] { runLength, zigzag[j] });
					runLength = 0;
				} else {
					runLength++;
				}
			}

		}

		// AC writing
		// O(n)
		for (int i = 0; i < YACList.size(); i++) {
			Entropy.writeAC(YACList.get(i)[0], YACList.get(i)[1]);
		}
		for (int i = 0; i < UACList.size(); i++) {
			Entropy.writeAC(UACList.get(i)[0], UACList.get(i)[1]);
		}
		for (int i = 0; i < VACList.size(); i++) {
			Entropy.writeAC(VACList.get(i)[0], VACList.get(i)[1]);
		}

		// Writing the whole thing
		SZLReaderWriter.writeSZLFile(destination, height, width, qualityFactor);
	}

}
